#!/bin/tcsh -ef
#############################################################################
#                                                                           #
# Title: Generate crystal averages                               					#
#                                                                           #
# (C) 2dx.org, GNU Plublic License.                                         #
#                                                                           #
# Created..........: 29/07/2016                                             #
# Last Modification: 29/07/2016                                             #
# Author...........: 2dx.org                                                #
#                                                                           #
#############################################################################
#
# SORTORDER: 3
#
#
# MANUAL: This script will calculate the average of the particles for every crystal. This can be useful to assess whether the particle extraction parameters used in the previous step are good. Also, these averages, which are analogous to 2D class averages in conventional Single Particle Analysis, are used in the next step for the first refinement in single-particle mode. This averaging procedure is similar to Correlation Averaging. Optionl Fourier Ring Correlations can be calculated for every crystal average to assess its quality.
#
# DISPLAY: Thread_Number
# DISPLAY: SPR_DIR
# DISPLAY: SPR_WhichStack
# DISPLAY: SPR_DoFRC
# DISPLAY: SPR_SigmaNorm
# DISPLAY: SPR_RadNorm
# DISPLAY: SPR_FRCThreshold
# DISPLAY: SPR_ShuffleParticleOrder
# DISPLAY: SPR_FilterType
# DISPLAY: SPR_ResCutoff
# DISPLAY: sample_pixel
#
#$end_local_vars
#
echo "<<@progress: 0>>"

set Thread_Number = ""
set bin_2dx = ""
set proc_2dx = ""

set scriptname = GenerateCrystalAverages

#
set SPR_DIR = ""
set SPR_WhichStack = ""
set SPR_DoFRC = ""
set SPR_SigmaNorm = ""
set SPR_RadNorm = ""
set sample_pixel = ""
set SPR_FRCThreshold = ""
set SPR_ShuffleParticleOrder = ""
set SPR_FilterType = ""
set SPR_ResCutoff  = ""
#
set SPR_STACKS_DIR = ${SPR_DIR}/stacks/
set SPR_STACK_ROOTNAME = 'particles'
set SPR_FRC_DIR = ${SPR_DIR}/FRC/
set scriptname = GenerateCrystalAverages
#$end_vars

set ccp4_setup = 'y'
source ${proc_2dx}/initialize
#

if ( -e LOGS/${scriptname}.results ) then
	mv LOGS/${scriptname}.results LOGS/${scriptname}.results.old
endif

if ( ${SPR_DoFRC} == "y" ) then
	mkdir -p ${SPR_FRC_DIR}
endif

if ( -e ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par ) then
	echo "::There is already a ${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par file in ${SPR_STACKS_DIR}! Will back it up."
	mv ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par.bkp
endif

if ( -e ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs ) then
	echo "::There is already a ${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs file in ${SPR_STACKS_DIR}! Will back it up."
	mv ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs.bkp
endif

# Let's make sure we don't use more threads than there are images to average particles from:
set n_img = `awk '{print $8}' ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_1_r1.par | sort | uniq | wc -l`
if ( ${n_img} < ${Thread_Number} ) then
	set Thread_Number = ${n_img}
else
# Also ensure that we don't have a wrong number of threads that will break parallelism:
	set load_img=`echo "${n_img}/${Thread_Number}" | bc -l`
	set load_img=`printf %.0f ${load_img}`
	while ( `echo "(${Thread_Number}-1)*${load_img} > ${n_img}" | bc -l` )
		@ Thread_Number--
		set load_img=`echo "${n_img}/${Thread_Number}" | bc -l`
		set load_img=`printf %.0f ${load_img}`
		# echo ${load_img}
	end
endif
echo "::Will use ${Thread_Number} threads for optimal load balancing."

if ( ${SPR_FilterType} == "FRC_weights" ) then
	set SPR_FilterType = "FRC"
endif
if ( ${SPR_FilterType} == "FRC^2" ) then
	set SPR_FilterType = "FRC2"
endif

set t = 1
set pids = ""
while ( $t <= ${Thread_Number} )
	${app_python} ${proc_2dx}/SPR_AverageParticles.py ${SPR_STACKS_DIR} ${SPR_STACK_ROOTNAME} ${SPR_WhichStack} ${SPR_DoFRC} ${SPR_FRC_DIR} ${SPR_SigmaNorm} ${SPR_RadNorm} ${sample_pixel} ${SPR_FRCThreshold} ${SPR_ShuffleParticleOrder} ${SPR_FilterType} ${SPR_ResCutoff} ${Thread_Number} ${t} &
	set pids = "$pids $!"
	@ t++
end

# Wait for all jobs to finish
while ( `ps -p ${pids} | wc -l` > 1 )
	wait
end

if ( -e ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg-0001.mrcs ) then
echo "::Creating single ${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs stack..."

	set t = 1

	# echo "::Creating single ${SPR_STACK_ROOTNAME}${i}.hdf stack..."
	while ( $t <= ${Thread_Number} )

		set num = `printf "%.4d" $t`

		# Append partial .mrcs file to the main one...
		if ( ${num} == "0001" ) then
			tail -c +0 ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg-${num}.mrcs >> ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs
		else
			tail -c +1025 ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg-${num}.mrcs >> ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs
		endif

		cat ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1-${num}.par >> ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par

		@ t++

	end
	echo "::Cleaning up partial ${SPR_STACK_ROOTNAME} stacks..."
	rm ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1-*.par
	rm ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg-*.mrcs

	${app_python} ${proc_2dx}/MRCheaderUpdate.py ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs

echo "::Done!"

awk '{$1=NR; print}' ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par > ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par.tmp
sed 's/ /    /g' ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par.tmp > ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par
# \mv ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par.tmp ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg_1_r1.par

echo "# IMAGE-IMPORTANT MRCS: ${SPR_STACKS_DIR}/${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs <STACK: ${SPR_STACK_ROOTNAME}_${SPR_WhichStack}_crystal-avg.mrcs>" >> LOGS/${scriptname}.results

if ( ${SPR_DoFRC} == "y" ) then

	foreach i (${SPR_FRC_DIR}/*.png)
		echo "# IMAGE: $i <PNG: `basename $i`>" >> LOGS/${scriptname}.results
	end

endif

echo "<<@progress: 100>>"

echo ":: "
echo ":: Done!"
echo ":: "
